package com.stone.starter.core.concurrent;

import org.apache.commons.collections4.CollectionUtils;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.scheduling.SchedulingTaskExecutor;

import java.util.Collection;
import java.util.concurrent.TimeUnit;
import java.util.concurrent.atomic.AtomicInteger;
import java.util.function.Consumer;

/**
 * 并发执行器
 *
 * @author Mr_wenpan@163.com 2021/7/15 10:30
 */
public class ConcurrentExecutor<T> {

    private static final Logger LOG = LoggerFactory.getLogger(ConcurrentExecutor.class);
    private static final Long DEFAULT_WAIT_MILLIS = 200L;

    /**
     * 线程休眠时间
     */
    private final long sleepMillis;
    /**
     * 指定执行任务的线程个数
     */
    private final Integer taskCount;
    /**
     * 待处理的数据集合
     */
    private final Collection<T> processDataList;
    /**
     * 线程池
     */
    private final SchedulingTaskExecutor taskExecutor;
    /**
     * 正在执行任务的线程个数计数器
     */
    private final AtomicInteger taskCounter = new AtomicInteger(0);

    public ConcurrentExecutor(Integer taskCount, Collection<T> processDataList, SchedulingTaskExecutor taskExecutor) {
        this.taskCount = taskCount;
        this.taskExecutor = taskExecutor;
        this.processDataList = processDataList;
        sleepMillis = DEFAULT_WAIT_MILLIS;
    }

    public ConcurrentExecutor(long sleepMillis, Integer taskCount, Collection<T> processDataList, SchedulingTaskExecutor taskExecutor) {
        this.sleepMillis = sleepMillis;
        this.taskCount = taskCount;
        this.processDataList = processDataList;
        this.taskExecutor = taskExecutor;
    }

    /**
     * 执行线程任务
     *
     * @param consumer 消费者接口
     * @author Mr_wenpan@163.com 2021/7/15 10:37 上午
     */
    public void invoke(Consumer<T> consumer) {

        if (CollectionUtils.isEmpty(processDataList)) {
            return;
        }

        // 分发任务
        for (T processData : processDataList) {
            distributeTask(consumer, processData);
        }

        // 等待所有任务执行完成
        while (taskCounter.get() > 0) {
            taskWait();
        }
    }

    /**
     * 分发任务
     *
     * @param consumer    消费者接口
     * @param processData 每个线程处理的数据
     * @author Mr_wenpan@163.com 2021/7/15 10:38 上午
     */
    private void distributeTask(Consumer<T> consumer, T processData) {

        taskCounter.incrementAndGet();

        // 如果正在执行任务的线程个数大于了用户指定的线程个数，则等待
        while (taskCounter.get() > taskCount) {
            taskWait();
        }
        taskExecutor.execute(() -> {
            try {
                consumer.accept(processData);
            } finally {
                taskCounter.decrementAndGet();
            }
        });
    }

    private void taskWait() {
        try {
            TimeUnit.MILLISECONDS.sleep(sleepMillis);
        } catch (Exception ex) {
            LOG.error("task wait error:", ex);
        }
    }
}